import os
import email
import pandas as pd
from tqdm import tqdm
from email.parser import Parser
from email.policy import default
from cot_utils import generate_cot

def parse_email_content(file_path):
    """解析邮件内容，返回主体文本"""
    try:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
        
        # 使用email模块解析
        email_parser = Parser(policy=default)
        msg = email_parser.parsestr(content)
        
        # 提取主题
        subject = msg.get('subject', '')
        
        # 提取正文
        body = ""
        if msg.is_multipart():
            for part in msg.walk():
                if part.get_content_type() == "text/plain":
                    try:
                        part_body = part.get_payload(decode=True).decode('utf-8', errors='ignore')
                        body += part_body + "\n"
                    except:
                        continue
        else:
            try:
                body = msg.get_payload(decode=True).decode('utf-8', errors='ignore')
            except:
                body = msg.get_payload()
        
        # 合并主题和正文
        full_text = f"Subject: {subject}\n\n{body}"
        # 清理文本
        full_text = full_text.replace('\r\n', ' ').replace('\n', ' ').strip()
        return full_text[:1000]  # 限制长度
    except Exception as e:
        print(f"Error processing {file_path}: {str(e)}")
        return ""

def process_directory(directory, label):
    """处理目录下的所有邮件"""
    emails = []
    labels = []
    
    for root, _, files in tqdm(list(os.walk(directory)), desc=f"Processing {label} emails"):
        for file in files:
            if file.startswith('.'):  # 跳过隐藏文件
                continue
            file_path = os.path.join(root, file)
            content = parse_email_content(file_path)
            if content:  # 只添加非空内容
                emails.append(content)
                labels.append(label)
    
    return emails, labels

# 处理ham邮件
print("处理正常邮件...")
ham_dir = "data/enron/maildir"  # Enron数据集的正常邮件目录
ham_emails, ham_labels = process_directory(ham_dir, "ham")

# 处理spam邮件
print("处理垃圾邮件...")
spam_dir = "data/enron/spam"  # Enron数据集的垃圾邮件目录
spam_emails, spam_labels = process_directory(spam_dir, "spam")

# 合并数据
emails = ham_emails + spam_emails
labels = ham_labels + spam_labels

# 创建DataFrame
df = pd.DataFrame({
    'text': emails,
    'label': labels
})

# 生成思维链提示
print("生成思维链提示...")
df['cot'] = df.apply(lambda row: generate_cot(row['text'], row['label']), axis=1)

# 显示基本信息
print("\n数据集信息:")
print(df['label'].value_counts())
print("\n示例:")
print(df.head())

# 保存为CSV
output_file = 'data/enron_processed.csv'
df.to_csv(output_file, index=False)
print(f"\n数据已保存至: {output_file}")

# 输出一些统计信息
print("\n数据集统计:")
print(f"总邮件数: {len(df)}")
print(f"正常邮件数: {len(df[df['label']=='ham'])}")
print(f"垃圾邮件数: {len(df[df['label']=='spam'])}")
print(f"平均文本长度: {df['text'].str.len().mean():.2f}字符") 